### 步骤
# 1.导入PDF文件
# 2.从PDF文件中提取文本并将其拆分为chunk
# 3.将这些小块发送到embedding 嵌入模型
# 4.将生成的embeddings嵌入保存到向量数据库
# 5.在向量数据库上执行相似性搜索以查找相似文档
# 6.检索相似文档并将结果反馈给用户

## 安装所需的软件包：pip install -r requirements.txt

## 1. 导入所需的库
from langchain.document_loaders import PyPDFLoader
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.output_parsers import StrOutputParser

from langchain_ollama import ChatOllama

from langchain_core.runnables import RunnablePassthrough
from langchain.retrievers.multi_query import MultiQueryRetriever

## 2. 设置文档路径和模型
doc_path = "../data/《中华人民共和国义务教育法》.pdf"
# model = "llama3.2"
model = "deepseek-r1:1.5b"
# model = "gemma:2b"
# model = "llava:7b"



## 3. 本地PDF文件加载
if doc_path:
    loader = PyPDFLoader(file_path=doc_path)
    data = loader.load()
    print("PDF文档加载完成....")
else:
    print("请上传PDF文件")

    # 预览第一页内容,打印前100个字符
content = data[0].page_content
print(content[:100])


# ==== 结束PDF加载 ====

# ==== 从PDF文件中提取文本并拆分成小块(Chunks) ====
# 拆分和分块,每个块的大小为1200字符，重叠部分为300字符
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=300)
chunks = text_splitter.split_documents(data)
print("文档 split 完成....")

print(f"chunks总数: {len(chunks)}")
# print(f"chunk[0]示例: {chunks[0]}")

# ===== 添加到向量数据库 ====
import ollama

# ollama.pull("nomic-embed-text")

vector_db = Chroma.from_documents(
    documents=chunks,
    embedding=OllamaEmbeddings(model="nomic-embed-text"), #使用 nomic-embed-text 模型生成文本嵌入
    collection_name="simple-rag", #存储在名为simple-rag的集合中
)
print("添加到向量数据库完成....")


## === 检索 ===

# 设置模型
llm = ChatOllama(model=model)


# 使用多重查询检索技术，根据用户问题生成多个查询，从向量数据库中检索相关文档
QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""你是一个人工智能语言模型助手，你的任务是生成五个不同版本的用户问题，
    以便从向量数据库中检索相关文档。通过对用户问题的多角度改写，你可以帮助用户克服基于距离的相似性搜索的局限性。
    请提供这些不同版本的问题，并使用换行符分隔。
    原始问题: {question}""",
)

retriever = MultiQueryRetriever.from_llm(
    vector_db.as_retriever(), llm, prompt=QUERY_PROMPT
)


# RAG 提示模板
# {context} 是一个占位符，用于在运行时插入特定的上下文信息。
# 在 LangChain 框架中，提示模板（PromptTemplate）允许通过占位符来动态地构建提示，以适应不同的输入和任务需求。
# 在实际应用中，{context} 通常由检索器（retriever）从向量数据库或其他数据源中获取的相关文档或信息填充。
# 当用户提出问题时，系统会检索与该问题相关的内容，并将这些内容作为 context 插入到提示模板中。
# 这样，语言模型就能够基于提供的上下文来生成更准确和相关的回答。
template = """请仅根据以下提供的上下文回答问题：
{context}
问题: {question}
"""

# 生成提示的对象
prompt = ChatPromptTemplate.from_template(template)

# 管道操作符（|），用于实现函数或对象的链式调用，参数解释如下：
# retriever:检索向量数据库，获取相关的文本内容
# RunnablePassthrough:直接传递用户输入的问题
# prompt:查询提示
# llm:使用之前定义的 ChatOllama 模型生成回答
# StrOutputParser：解析模型返回的字符串，并将其格式化为最终的输出
chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# 提问示例
# res = chain.invoke(input=("你是谁？"))
# res = chain.invoke(input=("你现在拥有什么知识？"))
# res = chain.invoke(input=("第二十三条内容是什么？"))
# print(f"模型: {res}")

# 实时对话循环
print("请输入您的问题（输入 '退出' 结束对话）：")
while True:
    user_input = input("用户: ")
    if user_input.lower() in ['退出', 'exit']:
        print("结束对话。")
        break
    response = chain.invoke(input=user_input)
    print(f"模型: {response}")